import random
import math

# Miller-Rabin素数测试
def is_prime(n, k=5):
    if n <= 1:
        return False
    if n <= 3:
        return True
    if n % 2 == 0:
        return False

    r, s = 0, n - 1
    while s % 2 == 0:
        r += 1
        s //= 2

    for _ in range(k):
        a = random.randrange(2, n - 1)
        x = pow(a, s, n)
        if x == 1 or x == n - 1:
            continue
        for _ in range(r - 1):
            x = pow(x, 2, n)
            if x == n - 1:
                break
        else:
            return False
    return True

# 生成大素数p和q
def generate_prime(bits):
    while True:
        p = random.getrandbits(bits)
        if is_prime(p):
            return p

def gcd(a, b):
    while b != 0:
        a, b = b, a % b
    return a

# 找到满足条件的e和d
def generate_keypair(p, q):
    n = p * q
    phi = (p-1) * (q-1)
    while True:
        e = random.randrange(2, phi)
        if gcd(e, phi) == 1:
            break
    d = pow(e, -1, phi)
    return ((e, n), (d, n))

# 加密
def encrypt(public_key, message):
    e, n = public_key
    encrypted_msg = [pow(ord(char), e, n) for char in message]
    return encrypted_msg

# 解密
def decrypt(private_key, encrypted_msg):
    d, n = private_key
    decrypted_msg = [chr(pow(char, d, n)) for char in encrypted_msg]
    return ''.join(decrypted_msg)

# 测试
if __name__ == "__main__":
    # 生成大素数p和q
    p = generate_prime(50)
    q = generate_prime(50)

    # 生成公钥和私钥
    public_key, private_key = generate_keypair(p, q)
    print(f"Public Key: {public_key}")
    print(f"Private Key: {private_key}")
    print(f"Prime p: {p}")
    print(f"Prime q: {q}")

    # 加密
    message = "Hello, world!"
    print(f"Original message: {message}")
    encrypted_msg = encrypt(public_key, message)
    print(f"Encrypted message: {encrypted_msg}")

    # 解密
    decrypted_msg = decrypt(private_key, encrypted_msg)
    print(f"Decrypted message: {decrypted_msg}")
